package com.kitten.chatglmsdk.executor.aigc;

import com.alibaba.fastjson.JSON;
import com.kitten.chatglmsdk.executor.Executor;
import com.kitten.chatglmsdk.model.*;
import com.kitten.chatglmsdk.session.Configuration;
import lombok.extern.slf4j.Slf4j;
import okhttp3.*;
import okhttp3.sse.EventSource;
import okhttp3.sse.EventSourceListener;
import org.jetbrains.annotations.Nullable;

import java.io.IOException;
import java.util.concurrent.CompletableFuture;

@Slf4j
public class GLMOldExecutor implements Executor {
    private final Configuration configuration;
    private final EventSource.Factory factory;

    public GLMOldExecutor(Configuration configuration) {
        this.configuration = configuration;
        this.factory = configuration.createRequestFactory();
    }
    public EventSource completions(ChatCompletionRequest chatCompletionRequest, EventSourceListener eventSourceListener) throws Exception {
        Request request = (new Request.Builder()).url(this.configuration.getApiHost().concat("api/paas/v3/model-api/{model}/sse-invoke").replace("{model}", chatCompletionRequest.getModel().getCode())).post(RequestBody.create(MediaType.parse("application/json"), chatCompletionRequest.toString())).build();
        return this.factory.newEventSource(request, eventSourceListener);
    }

    public CompletableFuture<String> completions(ChatCompletionRequest chatCompletionRequest) throws InterruptedException {
        final CompletableFuture<String> future = new CompletableFuture();
        final StringBuffer dataBuffer = new StringBuffer();
        Request request = (new Request.Builder()).url(this.configuration.getApiHost().concat("api/paas/v3/model-api/{model}/sse-invoke").replace("{model}", chatCompletionRequest.getModel().getCode())).post(RequestBody.create(MediaType.parse("application/json"), chatCompletionRequest.toString())).build();
        this.factory.newEventSource(request, new EventSourceListener() {
            public void onEvent(EventSource eventSource, @Nullable String id, @Nullable String type, String data) {
                ChatCompletionResponse response = (ChatCompletionResponse)JSON.parseObject(data, ChatCompletionResponse.class);
                if (EventType.add.getCode().equals(type)) {
                    dataBuffer.append(response.getData());
                } else if (EventType.finish.getCode().equals(type)) {
                    future.complete(dataBuffer.toString());
                }

            }
            public void onClosed(EventSource eventSource) {
                future.completeExceptionally(new RuntimeException("Request closed before completion"));
            }

            public void onFailure(EventSource eventSource, @Nullable Throwable t, @Nullable Response response) {
                future.completeExceptionally(new RuntimeException("Request closed before completion"));
            }
        });
        return future;
    }

    public ChatCompletionSyncResponse completionsSync(ChatCompletionRequest chatCompletionRequest) throws IOException {
        Request request = (new Request.Builder()).url(this.configuration.getApiHost().concat("api/paas/v3/model-api/{model}/invoke").replace("{model}", chatCompletionRequest.getModel().getCode())).header("Accept", "application/json").post(RequestBody.create(MediaType.parse("application/json"), chatCompletionRequest.toString())).build();
        OkHttpClient okHttpClient = this.configuration.getOkHttpClient();
        Response response = okHttpClient.newCall(request).execute();
        if (!response.isSuccessful()) {
            new RuntimeException("Request failed");
        }
        String string = response.body().string();
        log.info("测试结果：{}", string);
        return (ChatCompletionSyncResponse) JSON.parseObject(string, ChatCompletionSyncResponse.class);
    }

    @Override
    public ImageCompletionResponse imageGeneration(ImageCompletionRequest request) {
        throw new RuntimeException("旧版本模型不支持图像生成");
    }
}
